import json
import os
import sys
from pathlib import Path

import torch

from chinese_clip.model import CLIP


def load_chinese_clip(device, checkpoint_path, visual_name, text_name="RoBERTa-wwm-ext-base-chinese"):
    config_path = os.path.join(Path.cwd().parent, os.path.join('chinese_clip', 'model_configs'))
    visual_config_path = os.path.join(config_path, ''.join([visual_name, '.json']))
    text_config_path = os.path.join(config_path, ''.join([text_name, '.json']))

    with open(visual_config_path, 'r') as vf, open(text_config_path, 'r') as tf:
        visual_config = json.load(vf)
        text_config = json.load(tf)

        config_dict = {**visual_config, **text_config}

    model = CLIP(**config_dict)

    state_dict = torch.load(checkpoint_path, map_location="cpu")['state_dict']

    if next(iter(state_dict.items()))[0].startswith('module'):
        state_dict = {k[len('module.'):]: v for k, v in state_dict.items() if "bert.pooler" not in k}

    model.load_state_dict(state_dict)

    return model.to(device)


if __name__ == "__main__":
    load_chinese_clip('cuda', "E:\checkpoint\CN_CLIP\clip_cn_vit-b-16.pt", "ViT-B-16")













